# 这段代码使用了一些语言模型（LLM）和向量存储库（vector store）的功能，包括 DashScope Embeddings、Tongyi LLM 和 FAISS 向量存储库。
# 该代码使用了一种称为“管道”的技术，该技术允许将多个转换应用于单个输入，从而实现数据处理的高效和简洁。
# 该代码首先使用 Tongyi LLM 对输入文本进行编码，然后将其存储在 FAISS 向量存储库中。
# 然后，它使用 ChatPromptTemplate 对象来创建提示模板，该对象允许在上下文中添加问题。
# 最后，它使用管道将提示模板和 LLM 与向量存储库结合起来，以实现对输入问题的回答。
from _operator import itemgetter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.llms.tongyi import Tongyi
from langchain_community.vectorstores import FAISS
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

# 创建 LLM 对象和向量存储库对象
llm = Tongyi()
db = FAISS.from_texts(['张三工作在北京，他的电话号码是13812345678。'], embedding=DashScopeEmbeddings())

# 创建检索器
retriever = db.as_retriever()

# 创建提示模板
template = """
根据下面内容，回答问题,使用{language}语言回答：
{context}

问题：{question}
"""
prompt = ChatPromptTemplate.from_template(template)

# 创建管道
chain = (
        {
            "context": itemgetter('question') | retriever,
            "question": itemgetter('question'),
            'language': itemgetter('language')
        }
        | prompt
        | llm
        | StrOutputParser()
)

# 调用管道，并提供输入
res = chain.invoke({'question': '张三的电话号码是多少？在哪里工作？', 'language': '中文'})

# 打印输出
print(res)

# 再次调用管道，并提供相同的输入
template = """
根据下面内容，回答问题,使用{language}语言回答：
{context}

问题：{question}
"""
prompt = ChatPromptTemplate.from_template(template)

chain = (
        {
            "context": itemgetter('question') | retriever,
            "question": itemgetter('question'),
            'language': itemgetter('language')
        }
        | prompt
        | llm
        | StrOutputParser()
)

res = chain.invoke({'question': '张三的电话号码是多少？在哪里工作？', 'language': '中文'})
print(res)

